#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Software: PyCharm
# @Version : Python-
# @Author  : Shengji He
# @Email   : hsjbit@163.com
# @File    : ConvertBST2GreaterTree.py
# @Time    : 2020/9/21 10:21
# @Description:


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


class Solution:
    def convertBST(self, root: TreeNode) -> TreeNode:
        """
        Given a Binary Search Tree (BST), convert it to a Greater Tree such that every key of the original BST is
        changed to the original key plus sum of all keys greater than the original key in BST.
        Example:
            Input: The root of a Binary Search Tree like this:
                          5
                        /   \
                       2     13
            Output: The root of a Greater Tree like this:
                         18
                        /   \
                      20     13
        Note: This question is the same as 1038: https://leetcode.com/problems/binary-search-tree-to-greater-sum-tree/

        :param root:
        :return:
        """
        def dfs(root: TreeNode):
            nonlocal total
            if root:
                dfs(root.right)
                total += root.val
                root.val = total
                dfs(root.left)
        total = 0
        dfs(root)
        return root

    def convertBST2(self, root: TreeNode) -> TreeNode:
        def getSuccessor(node: TreeNode) -> TreeNode:
            succ = node.right
            while succ.left and succ.left != node:
                succ = succ.left
            return succ

        total = 0
        node = root
        while node:
            if not node.right:
                total += node.val
                node.val = total
                node = node.left
            else:
                succ = getSuccessor(node)
                if not succ.left:
                    succ.left = node
                    node = node.right
                else:
                    succ.left = None
                    total += node.val
                    node.val = total
                    node = node.left
        return root



if __name__ == '__main__':
    S = Solution()
    root = TreeNode(5)
    root.left = TreeNode(2)
    root.right = TreeNode(13)
    # print(S.convertBST(root))
    print(S.convertBST2(root))
    print('done')
